package nil.ed.benchmark.stomp;

import co.paralleluniverse.fibers.Fiber;
import nil.ed.benchmark.common.concurrent.SleeperLatch;
import nil.ed.benchmark.stomp.bean.SessionBean;
import org.apache.commons.lang3.StringUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.springframework.messaging.converter.*;
import org.springframework.messaging.simp.stomp.*;
import org.springframework.web.socket.WebSocketHttpHeaders;
import org.springframework.web.socket.client.WebSocketClient;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;
import org.springframework.web.socket.messaging.WebSocketStompClient;
import org.springframework.web.socket.sockjs.client.SockJsClient;
import org.springframework.web.socket.sockjs.client.WebSocketTransport;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.IntStream;

/**
 * @author lidelin
 */
public class StompAutoSender extends StompSessionHandlerAdapter{

    private static Logger log = LogManager.getLogger(StompAutoSender.class);

    private static Map<String, CustomStatisticMessageHandler> handlerMap = new ConcurrentHashMap<>();

    private List<SessionBean> sessionList;

    private AtomicInteger failedCount = new AtomicInteger(0);

    private static SleeperLatch latch = new SleeperLatch();

    public static void main(String[] args) throws Exception{
        StompAutoSender autoSender = new StompAutoSender();
        int count = 200;
        autoSender.sessionList = Collections.synchronizedList(new ArrayList<>(count));
        IntStream.range(0, count)
                .parallel()
                .forEach(i -> {
                    autoSender.initOneClient(Collections.singletonList(autoSender));
                });
        int failedCount = autoSender.failedCount.get();
        log.info("Total count = {}, succeed count = {}, failed count = {}", count, count - failedCount, failedCount);
        sleepNonException(5000);
        long maxTime = 1000;
        long messageCount = 1000;
        long loopCount = messageCount;
        AtomicInteger failedMessageCount = new AtomicInteger(0);
        List<Fiber<Object>> fibers = new LinkedList<>();
        while (loopCount-- > 0) {
            StompSession session = autoSender.randomSelect();
            Fiber<Object> fiber = new Fiber<Object>(() -> {
                try {
                    autoSender.asyncSendMessage(session, autoSender.getDestination(), "" + System.currentTimeMillis());
                } catch (Exception e) {
                    log.error("{}", e.getMessage());
                    failedMessageCount.incrementAndGet();
                }
            }).start();
            fibers.add(fiber);
        }

        for (Fiber<Object> f : fibers) {
            f.get();
        }

        latch.sleep(true);

        handlerMap.values().stream()
                .peek(CustomStatisticMessageHandler::statisticLatency)
                .forEach(h -> {
                    log.info("Max Latency = {}, Min Latency = {}, Avg Latency = {}, Message Count = {}, Error Count = {}",
                            h.getMaxLatency(), h.getMinLatency(), h.getAvgLatency(), h.getMessageCount(), h.getErrorCount());
                });

        log.info("Total Message Count = {}, Failed Message Count = {}, Failed/Total = {}",
                messageCount, failedMessageCount.get(), 1.0 * failedMessageCount.get() / messageCount);
    }

    @Override
    public void afterConnected(StompSession session, StompHeaders connectedHeaders) {
        super.afterConnected(session, connectedHeaders);
        CustomStatisticMessageHandler handler = new CustomStatisticMessageHandler();
        SessionBean sessionBean = new SessionBean(session);
        handlerMap.put("/topic/groups.1" + session.getSessionId(), handler);
        session.subscribe("/topic/echo.1", this);
        session.subscribe("/topic/group.1", handler);
        session.subscribe("/user/topic/oneToOne.1", this);
        sessionList.add(sessionBean);
    }

    public StompSession randomSelect() {
        int index = ThreadLocalRandom.current().nextInt(sessionList.size());
        return sessionList.get(index).getSession();
    }

    public String getDestination() {
        return "/topic/group.1";
    }

    private static boolean sleepNonException(long timeout) {
        try {
            Thread.sleep(timeout);
            return true;
        } catch (InterruptedException e) {
            return false;
        }
    }

    private static final byte[] EMPTY_STR_BYTES = "".getBytes();
    private void asyncSendMessage(StompSession session, String dest, String message) {
        byte[] payload = Optional.ofNullable(message)
                                .filter(StringUtils::isNotBlank)
                                .map(String::getBytes)
                                .orElse(EMPTY_STR_BYTES);
        StompHeaders headers = new StompHeaders();
        headers.setDestination(dest);
        headers.set("ts", "" + System.currentTimeMillis());
        session.send(headers, payload);
    }

    public void initOneClient(List<StompSessionHandler> handlers) {
        Collection<MessageConverter> converters = new LinkedList<>();
        Collections.addAll(converters,
                new StringMessageConverter(),
                new ByteArrayMessageConverter());
        MessageConverter messageConverter = new CompositeMessageConverter(converters);
        WebSocketClient webSocketClient = new StandardWebSocketClient();
        SockJsClient client = new SockJsClient(Collections.singletonList(new WebSocketTransport(webSocketClient)));
        WebSocketStompClient webSocketStompClient = new WebSocketStompClient(client);
        WebSocketHttpHeaders headers = new WebSocketHttpHeaders();
        headers.add("Cookie", "jwplayer.volume=0; JSESSIONID=YWEwZGJmOTAtN2QyYi00YTY5LThkMmEtOWZlMGUxNjkxMjdk; SESSION=Yzk3MTMwZmUtNWUyNS00ZWY4LThiY2UtZTc4Y2Q2NjA3NWI4");
        StompHeaders stompHeaders = new StompHeaders();
        stompHeaders.add("room-id", "1");
        webSocketStompClient.setMessageConverter(messageConverter);
        webSocketStompClient.connect("ws://localhost:12011/chat/room", headers, stompHeaders, this);
    }

    @Override
    public void handleException(StompSession session, StompCommand command, StompHeaders headers, byte[] payload, Throwable exception) {
        super.handleException(session, command, headers, payload, exception);
        log.error("", exception);
    }

    public static class CustomStatisticMessageHandler extends StompSessionHandlerAdapter {

        private AtomicInteger messageCounter = new AtomicInteger(0);
        private AtomicInteger errorCounter = new AtomicInteger(0);
        private BlockingQueue<Long> latencyQueue = new LinkedBlockingQueue<>();
        private long maxLatency;
        private long minLatency;
        private long avgLatency;

        @Override
        public void handleException(StompSession session, StompCommand command, StompHeaders headers, byte[] payload, Throwable exception) {
            super.handleException(session, command, headers, payload, exception);
            errorCounter.incrementAndGet();
        }

        @Override
        public void handleFrame(StompHeaders headers, Object payload) {
            messageCounter.incrementAndGet();
            latencyQueue.offer(calLatency(headers, payload));
            latch.register(5000L);
            latch.setNotifyFlag(true);
        }

        public long calLatency(StompHeaders headers, Object payload) {
            long sendTime = Long.parseLong("" + headers.getFirst("ts"));
            return System.currentTimeMillis() - sendTime;
        }

        public void statisticLatency() {
            long maxLatency = -1L;
            long minLatency = Long.MAX_VALUE;
            long sumLatency = -1L;
            long size = latencyQueue.size();
            Long latency;
            while((latency = latencyQueue.poll()) != null) {
                maxLatency = Math.max(maxLatency, latency);
                minLatency = Math.min(minLatency, latency);
                sumLatency += latency;
            }
            this.maxLatency = maxLatency;
            this.minLatency = minLatency;
            this.avgLatency = size == 0 ? -1 : sumLatency / size;
        }

        public long getMaxLatency() {
            return maxLatency;
        }

        public long getMinLatency() {
            return minLatency;
        }

        public long getAvgLatency() {
            return avgLatency;
        }

        public int getMessageCount() {
            return messageCounter.get();
        }

        public int getErrorCount() {
            return errorCounter.get();
        }

    }

}
